# -*- coding: utf-8 -*-
# @Time    : 2021/9/29 9:27
# @Author  : huangwei
# @File    : split_table.py
# @Software: PyCharm
import cv2
import os
import layoutparser as lp
from method import create_dir


def split_table( filepath, table_layout ):
    img = cv2.imread(filepath)
    height, width = img.shape[0:2]
    layout_res = table_layout.detect(img[..., ::-1])

    # 创建输出文件夹
    tmp_file_name = os.path.basename(filepath)
    file_name, _ = os.path.splitext(tmp_file_name)
    output_dir = "output/{}".format(file_name)
    create_dir(output_dir)

    # 显示结果
    # show_img = lp.draw_box(img, layout_res, box_width=3, show_element_type=True)
    # show_img.show()

    # 根据结果裁剪出需要识别的表格
    for region in layout_res:
        region_type = region.type
        if region_type == 'Table':
            rect_data = region.block
            x1 = int(rect_data.x_1)
            y1 = int(rect_data.y_1)
            x2 = int(rect_data.x_2)
            y2 = int(rect_data.y_2)

            # 边缘向外扩展 10
            x1 = max(0, x1 - 10)
            y1 = max(0, y1 - 10)
            x2 = min(width, x2 + 10)
            y2 = min(height, y2 + 10)

            new_img = img[y1:y2, x1:x2]

            # 根据截取的图片的中间位置的y坐标命名
            center_y = int((y1 + y2) / 2)
            temp_img_path = "{0}/{1}.png".format(output_dir, center_y)

            cv2.imwrite(temp_img_path, new_img)


img_path = "images/resource2.png"

# 加载识别表格的模型文件
config_path = "lp://TableBank/ppyolov2_r50vd_dcn_365e_tableBank_word/config"
model_path = "inference/ppyolov2_r50vd_dcn_365e_tableBank_word"
table_layout = lp.PaddleDetectionLayoutModel(config_path=config_path, threshold=0.1, model_path=model_path,
                                             thread_num=5)

split_table(img_path, table_layout)
